'''
Created on Jul 21, 2011

@author: anupams
'''

from com.anupam.graphs.GraphsAnupam import SampleGraphObject
import copy

class KrusKalMST(object):
  """
    Class for finding Minimum Spanning Tress, Using Kruskal's algorithm.
  """

  def __init__(self, graphDict, nodes):
    """
    Default constructor of the class.
    
    Args:
      graphDict : Adjucancy list of the graph
      nodes : graph node objects @see{ColoredGraphNode}
    """
    self.graphDict = graphDict
    self.nodes = nodes
    self.sortedEgdeByWeight = []
    ## Keeps a count of edges in the graph.
    self.edgeCount = 0
    ## Holds the minimum spanning tree of graph.
    self.MSTEdges = []
    self.tempEdgeDict = {}
    ## Flag to indicate we have detected a cycle for the edge being inserted to the minSpanning Tree.
    self.cycleDetected = False
    self.totalTreeWeight = 0
    # Intialize with empty array all the nodes. This array will hold only the edges that are part 
    # of MST from this particular node. i.e. for each key create an array containing the other 
    # end of edge that forms MST
    for key in self.nodes:
      self.tempEdgeDict[key] = []

  def __deleteFromDict(self, smallestEdge, parentOfSmallest):
    """ 
    Deletes those edges from dict that are in correct sorted order.
    Therfore preventing re-traversal of smallestEdge.
    """
    childNodes = self.graphDict[parentOfSmallest]
    childNodes.remove(smallestEdge)

  def __checkAndInsert(self, edgeToInsert):
    """
      Inserts an edge in sortedEdge list only for cases that have no
      reverse entry in the list.
      i.e. If edge i -> g is already inserted then g -> i won't be inserted.
    """
    # TODO(anupam): Handle the case when there are multiple edges
    # between the nodes of different weights.
    for i in range(0, len(self.sortedEgdeByWeight)):
      currentEdgeParent = self.sortedEgdeByWeight[i]['parent']
      currentEdgeChildName = self.sortedEgdeByWeight[i]['childname']
      if ((currentEdgeParent == edgeToInsert['childname']) 
           and (currentEdgeChildName == edgeToInsert['parent'])):
        return
    self.sortedEgdeByWeight.append(edgeToInsert)

  def __calculateEdgesInGraph(self):
    """
      Counts total number of edges in the graph.
    """
    for node in self.graphDict:
      nodeChildren = self.graphDict[node]
      self.edgeCount = self.edgeCount + len(nodeChildren)
    return self.edgeCount

  def __findEdgeWithLeastWeight(self):
    """
      Finds an edge with the least weight in the graph.
      Once such edge is found it removes that edge from the graph.
      This brings down the comparisons to almost nlog(n) size.
    """
    smallestEdge = None
    parentOfSmallest = None

    for node in self.graphDict:
      nodeChildren = self.graphDict[node]
      for i in range(0, len(nodeChildren)):
        eachChild = nodeChildren[i]
        if smallestEdge == None:
          smallestEdge = eachChild
          parentOfSmallest = node
        elif eachChild['weight'] <= smallestEdge['weight']:
          smallestEdge = eachChild
          parentOfSmallest = node
    temp = copy.deepcopy(smallestEdge)
    temp['parent'] = parentOfSmallest

    ### Deletes the recently discovered smallest weighted edge from graph.
    self.__deleteFromDict(smallestEdge, parentOfSmallest)

    ## Prevents duplicate insertion of same edge f--> g and g--> f repersent
    ## same edge & only one of these should be inserted.
    self.__checkAndInsert(temp)

  def sortEdgesByWeight(self):
    """
    Sorts the edges in the graph according to edge weight.
    It's time complexity is approximately nlog(n), @see: {__findEdgeWithLeastWeight} 
    """
    for i in range(0, self.__calculateEdgesInGraph()):
      self.__findEdgeWithLeastWeight()

  
  def __detectCycle(self, originatingVertex, rightVertex, edgesEncountered):
    """
      originatingVertex : Left vertex of edge being investigated.This value is passed
                          as it is to furthur recursive calls to __detectCycle method.
      rightVertex : Right vertex of edge being investigated.
      edgesEncountered : A key value representation of edges that are part of MST.
                         key is the name of the originating vertex and value is a
                         list of ending vertex, Such that each combination of key and list
                         value is a edge in MST originating from that key. @attention: It's
                         not a exhastive list of edges from that originating vertex.
    """
    ## List of all the vertex that have a edge from rightVertex and are Part of MST.
    outwardVertexLst = self.tempEdgeDict.get(rightVertex)
    if not outwardVertexLst:
      return
    #print 'for rightVertex : ' + rightVertex

    for i in range(0, len(outwardVertexLst)):
      endingVertex = outwardVertexLst[i]
      if edgesEncountered.get(rightVertex + endingVertex):
        ## Already encountered this edge while traversing MST.Check the edge color.
        ## TODO(anupam) : Below if check looks redundant . Check if can be removed.
        if edgesEncountered.get(rightVertex + endingVertex) == 'Red':
          ## Red color indicates a cycle if the edge is added to MST.
          #print 'this edge is causing a cycle : ' + rightVertex + endingVertex
          ## Set the cycle detection flag.
          self.cycleDetected = True
          return
        else:
          allAdjacentNodes = self.graphDict.get(rightVertex)
          ## inspect all the edges in graph for rightVertex
          for k in range(0, len(allAdjacentNodes)):
            temp = edgesEncountered.get(rightVertex + allAdjacentNodes[k]['childname'])
            ## If any of the edge is red it means we have a cycle
            if temp:
              if temp == 'Red':
                self.cycleDetected = True
                return

        ## Grey edge means , coninue exploring the outwardVertexLst.
        ## Grey colored edge should only occour when an already discovered MST
        ## edge is accesed from opposite direction. i.e edge g --> h is discovered
        ## then the edge h ---> g should also be considered discovered and be grey in color.
        ## This helps in preventing infinite recursion.
        ## TODO (anupam) : check weather this if check is required.
        if edgesEncountered.get(rightVertex + endingVertex) == 'Grey':
          #print 'we are traversing already encountered edge : ' + rightVertex + endingVertex 
          continue
      else:
        edgesEncountered[rightVertex + endingVertex] = 'Grey'
        edgesEncountered[endingVertex + rightVertex] = 'Grey'
        if self.tempEdgeDict.get(endingVertex):
          ## Make a recursive call to furthur explore the currently formed MST for cycle.
          self.__detectCycle(originatingVertex, endingVertex, edgesEncountered)

  def __startProcessingEdge(self, edgeToBeChecked):
    """
      edgeToBeChecked : Edge taken from { @see: sortedEgdeByWeight } and it will checked
                        if inserting it to MST will cause any cycle.
    """
    leftVertext = edgeToBeChecked['parent']
    rightVertex = edgeToBeChecked['childname']
    ## Reset cycle detection flag for this edge.It's value will be set in @see: {__detectCycle}
    self.cycleDetected = False

    edgesEncountered = {}
    ## Insert the edge that is being investigated. Set it's color to red.
    ## Red colored edge if encountered while iterating through the edges in @see: {tempEdgeDict}
    ## indicates that edgeToBeChecked will cause a cycle in the MST.
    edgesEncountered[leftVertext + rightVertex] = 'Red'
    self.__detectCycle(leftVertext, rightVertex, edgesEncountered)
    return self.cycleDetected


  def generateMST(self):
    """
      Generates an MST from the provided graph structure.
    """
    ## Insert the edge with least weight intially.
    if not len(self.MSTEdges):
      currEdge = self.sortedEgdeByWeight[0]
      self.MSTEdges.append(currEdge)
      self.tempEdgeDict[currEdge['parent']].append(currEdge['childname'])
      self.tempEdgeDict[currEdge['childname']].append(currEdge['parent'])
      self.totalTreeWeight = self.totalTreeWeight + currEdge['weight']

    for i in range(1, len(self.sortedEgdeByWeight)):
      currEdge = self.sortedEgdeByWeight[i]
      ## print ' current state : ' + str(self.MSTEdges)
      ## print 'before add : ' + str(currEdge)

      if not self.__startProcessingEdge(currEdge):
        ## Insert the edge formed from left to right vertex
        self.tempEdgeDict[currEdge['parent']].append(currEdge['childname'])
        ## Insert the edge formed from right to left vertex
        self.tempEdgeDict[currEdge['childname']].append(currEdge['parent'])
        self.MSTEdges.append(currEdge)
        self.totalTreeWeight = self.totalTreeWeight + currEdge['weight']

if __name__ == '__main__':
  directedGraphSample = SampleGraphObject()
  graphDict = directedGraphSample.undirectedGraphKruskals()
  traverse = KrusKalMST(graphDict, directedGraphSample.nodes)
  traverse.sortEdgesByWeight()
  print 'Graph sorted according to edge weight : ' + str(traverse.sortedEgdeByWeight)
  
  ## Repopulate graph dictionary.
  traverse.graphDict = directedGraphSample.undirectedGraphKruskals()
  traverse.generateMST()
  
  ## print the output
  print 'Edges forming MST for the graph : ' + str(traverse.MSTEdges)
  print 'tree weight : ' + str(traverse.totalTreeWeight)

